import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import json
import sys
import os
import warnings
import math
from sklearn.feature_selection import f_classif
from scipy.stats import f_oneway
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def replace_invalid_floats(data):
    """Replace NaN, inf, and -inf with None."""
    if isinstance(data, dict):
        return {key: replace_invalid_floats(value) for key, value in data.items()}
    elif isinstance(data, list):
        return [replace_invalid_floats(item) for item in data]
    elif isinstance(data, float):
        if math.isnan(data) or math.isinf(data):
            return None
    return data

def calculate_pairwise_anova(df):
    """
    计算特征之间的方差分析F值矩阵
    """
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    n_features = len(numeric_cols)

    # 初始化F值矩阵
    f_matrix = pd.DataFrame(index=numeric_cols, columns=numeric_cols, dtype=float)

    for i, feature1 in enumerate(numeric_cols):
        for j, feature2 in enumerate(numeric_cols):
            if i == j:
                # 对角线设为最大值（自身比较）
                f_matrix.loc[feature1, feature2] = 1.0
            else:
                try:
                    # 将feature1作为分组变量，feature2作为观测变量
                    feature1_data = df[feature1].dropna()
                    feature2_data = df[feature2].dropna()

                    # 找到公共索引
                    common_idx = feature1_data.index.intersection(feature2_data.index)
                    if len(common_idx) < 10:
                        f_matrix.loc[feature1, feature2] = 0.0
                        continue

                    # 获取公共数据
                    f1_values = feature1_data.loc[common_idx]
                    f2_values = feature2_data.loc[common_idx]

                    # 将feature1离散化为分组变量
                    if f1_values.nunique() > 10:
                        # 如果唯一值太多，使用分位数分组
                        f1_categorical = pd.qcut(f1_values, q=4, duplicates='drop')
                    else:
                        # 如果唯一值较少，直接使用
                        f1_categorical = f1_values

                    # 按分组分离feature2的值
                    groups = []
                    for group_name in f1_categorical.unique():
                        if pd.isna(group_name):
                            continue
                        group_mask = (f1_categorical == group_name)
                        group_data = f2_values[group_mask]
                        if len(group_data) > 0:
                            groups.append(group_data.values)

                    # 计算F统计量
                    if len(groups) >= 2:
                        f_stat, p_val = f_oneway(*groups)
                        if not np.isnan(f_stat) and not np.isinf(f_stat):
                            # 标准化F值到0-1范围，便于热力图显示
                            normalized_f = min(f_stat / 100.0, 1.0)
                            f_matrix.loc[feature1, feature2] = normalized_f
                        else:
                            f_matrix.loc[feature1, feature2] = 0.0
                    else:
                        f_matrix.loc[feature1, feature2] = 0.0

                except Exception as e:
                    print(f"计算 {feature1} vs {feature2} 的F值时出错: {e}")
                    f_matrix.loc[feature1, feature2] = 0.0

    return f_matrix

def main():
    try:
        if len(sys.argv) != 3:
            result = {
                'status': 'error',
                'message': '参数错误，需要提供CSV文件路径和选择的特征JSON文件路径'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        csv_file_path = sys.argv[1]
        features_json_file = sys.argv[2]

        # 验证输入文件是否存在
        if not os.path.exists(csv_file_path):
            result = {
                'status': 'error',
                'message': f'CSV文件不存在: {csv_file_path}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        if not os.path.exists(features_json_file):
            result = {
                'status': 'error',
                'message': f'特征JSON文件不存在: {features_json_file}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 读取特征列表
        try:
            with open(features_json_file, 'r', encoding='utf-8') as f:
                selected_features = json.load(f)

            if not isinstance(selected_features, list):
                result = {
                    'status': 'error',
                    'message': f'特征JSON文件格式错误，应为列表格式'
                }
                print(json.dumps(result, ensure_ascii=False))
                return

        except json.JSONDecodeError as e:
            result = {
                'status': 'error',
                'message': f'特征JSON文件解析失败: {str(e)}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 读取数据
        df = pd.read_csv(csv_file_path)

        # 确保所有特征字段都存在
        missing_features = [f for f in selected_features if f not in df.columns]
        if missing_features:
            result = {
                'status': 'error',
                'message': f'数据中缺少字段: {missing_features}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 选择20个字段的数据（包含qyb）
        selected_df = df[selected_features].copy()

        # 检查数据有效性
        if selected_df.dropna().empty:
            result = {
                'status': 'error',
                'message': '选定的特征数据全部为缺失值'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 计算方差分析F值矩阵
        print("开始计算特征间的方差分析F值矩阵...")
        f_matrix = calculate_pairwise_anova(selected_df)

        # 处理无效值
        f_matrix_for_json = f_matrix.where(pd.notnull(f_matrix), None)
        f_matrix_for_json = f_matrix_for_json.applymap(
            lambda x: None if (isinstance(x, float) and (math.isnan(x) or math.isinf(x))) else x
        )

        # 转换为字典
        f_matrix_dict = f_matrix_for_json.to_dict()

        # 创建热力图
        heatmap_path = None
        try:
            plt.figure(figsize=(16, 14))

            # 生成热力图
            sns.heatmap(f_matrix,
                        annot=True,
                        cmap='YlOrRd',  # 使用黄-橙-红色谱，更适合F值显示
                        square=True,
                        fmt='.3f',
                        cbar_kws={"shrink": .8, "label": "标准化F值"},
                        linewidths=0.5)

            plt.title(f'{len(selected_features)}个重要特征方差分析F值热力图\n(包含QYB目标变量)',
                      fontsize=16, fontweight='bold', pad=20)
            plt.xlabel('特征 (作为观测变量)', fontsize=12)
            plt.ylabel('特征 (作为分组变量)', fontsize=12)
            plt.xticks(rotation=45, ha='right')
            plt.yticks(rotation=0)
            plt.tight_layout()

            # 确保输出目录存在
            os.makedirs('temp/output', exist_ok=True)
            heatmap_path = 'temp/output/anova_f_values_heatmap.png'
            plt.savefig(heatmap_path, dpi=300, bbox_inches='tight')
            plt.close()

            print(f"方差分析F值热力图已保存: {heatmap_path}")

        except Exception as e:
            print(f"热力图生成失败: {str(e)}", file=sys.stderr)

        # 分析高F值特征对
        high_f_pairs = []
        for i in range(len(f_matrix.columns)):
            for j in range(len(f_matrix.columns)):
                if i != j:  # 排除对角线
                    f_value = f_matrix.iloc[i, j]
                    if not (math.isnan(f_value) or math.isinf(f_value)) and f_value > 0.5:  # F值阈值
                        high_f_pairs.append({
                            'grouping_feature': f_matrix.index[i],
                            'observed_feature': f_matrix.columns[j],
                            'f_value': float(f_value)
                        })

        # 如果qyb在特征中，分析其他特征对qyb的F值
        qyb_f_values = {}
        if 'qyb' in f_matrix.columns:
            for col in f_matrix.columns:
                if col != 'qyb':
                    qyb_f_values[col] = f_matrix_for_json.loc[col, 'qyb']  # col作为分组变量，qyb作为观测变量

        # 构建结果
        result = {
            'status': 'success',
            'message': '第二次方差分析完成 - 20个特征F值热力图生成',
            'analysis_type': 'anova_f_values_heatmap',
            'analysis_method': 'ANOVA (方差分析)',
            'f_values_matrix': f_matrix_dict,
            'features_analyzed': selected_features,
            'matrix_shape': list(f_matrix.shape),
            'high_f_pairs': high_f_pairs,
            'high_f_pairs_count': len(high_f_pairs),
            'qyb_f_values': qyb_f_values,
            'data_info': {
                'total_rows': len(df),
                'valid_rows': len(selected_df.dropna()),
                'features_count': len(selected_features)
            },
            'interpretation': {
                'f_values_meaning': 'F值表示特征作为分组变量时对观测变量的区分能力',
                'high_f_threshold': 0.5,
                'matrix_explanation': '行为分组变量，列为观测变量'
            }
        }

        if heatmap_path:
            result['heatmap_path'] = heatmap_path

        # 过滤无效值并输出
        result = replace_invalid_floats(result)
        print(json.dumps(result, ensure_ascii=False))

    except Exception as e:
        error_result = {
            'status': 'error',
            'message': f'第二次方差分析失败: {str(e)}'
        }
        print(json.dumps(error_result, ensure_ascii=False))

if __name__ == "__main__":
    main()
